import torch
import torchcontrib
from tqdm import tqdm
import numpy as np
from torch.cuda import amp
from torchnet.meter import AverageValueMeter
import torch.nn as nn


class Epoch:
    def __init__(self, scaler, model, criterion, metric, stage, device, use_amp):
        self.scaler = scaler
        self.model = model
        self.criterion = criterion
        self.metric = metric
        self.device = device
        self.stage = stage
        self.use_amp = use_amp
        self._to_device()

    def _to_device(self):
        self.model.to(self.device)
        if isinstance(self.criterion, nn.Module):
            self.criterion.to(self.device)
        self.metric.to(self.device)

    def batch_update(self, x, y):
        raise NotImplementedError

    def on_epoch_start(self):
        pass

    def on_epoch_end(self, epoch, dataloader):
        pass

    def run(self, epoch, dataloader):
        self.on_epoch_start()

        loss_meter = AverageValueMeter()
        iou_meter = AverageValueMeter()
        recall_meter = AverageValueMeter()
        f1_meter = AverageValueMeter()
        precision_meter = AverageValueMeter()
        if self.stage == 'train':
            print(('\n' + '%10s' * 8) % (self.stage, 'steps', 'gpu', 'loss', 'precision', 'recall', 'f1', 'iou'))

        pbar = tqdm(enumerate(dataloader), total=len(dataloader))

        for step, (x, y) in pbar:
            if isinstance(x, list):
                for i in range(len(x)):
                    x[i] = x[i].to(self.device)
            else:
                x = x.to(self.device)
            if isinstance(y, list):
                for i in range(len(y)):
                    y[i] = y[i].to(self.device)
            else:
                y = y.to(self.device)

            loss, y_pred = self.batch_update(x, y)

            loss_value = loss.cpu().detach().numpy()
            metrics = self.metric(y_pred, y)
            if np.nan != loss_value:
                loss_meter.add(loss_value)
                precision_meter.add(metrics[0])
                recall_meter.add(metrics[1])
                f1_meter.add(metrics[2])
                iou_meter.add(metrics[3])

            mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0)  # (GB)
            s = ('%10s' * 3 + '%10.4g' * 5) % (self.stage, '%g/%g' % (step, len(dataloader) - 1),
                                               mem, loss_value, metrics[0], metrics[1], metrics[2], metrics[3])
            pbar.set_description(s)

        logs = {
            'loss': loss_meter.mean,
            'precision': precision_meter.mean,
            'recall': recall_meter.mean,
            'f1': f1_meter.mean,
            'iou': iou_meter.mean
        }

        self.on_epoch_end(epoch, dataloader)
        return logs


class TrainEpoch(Epoch):

    def __init__(self, scaler, model, criterion, metric, optimizer, stage, device='cuda', use_amp=True):
        super().__init__(scaler, model, criterion, metric, stage, device, use_amp)
        self.optimizer = optimizer

    def on_epoch_start(self):
        self.model.train()

    def batch_update(self, x, y):
        self.optimizer.zero_grad()
        if self.use_amp:
            with amp.autocast():
                pred = self.model.forward(x)
                if isinstance(self.criterion, nn.Module):
                    loss = self.criterion(pred, y)
                else:
                    if self.device == 'cuda':
                        loss = self.criterion(pred, y, per_image=False)
                    else:
                        loss = self.criterion(pred, y, per_image=True)
            self.scaler.scale(loss).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            pred = self.model.forward(x)
            if isinstance(self.criterion, nn.Module):
                loss = self.criterion(pred, y)
            else:
                if self.device == 'cuda':
                    loss = self.criterion(pred, y, per_image=False)
                else:
                    loss = self.criterion(pred, y, per_image=True)
            loss.backward()
            self.optimizer.step()

        if isinstance(pred, list):
            pred = pred[0]
        return loss, pred


class ValEpoch(Epoch):

    def __init__(self, scaler, model, criterion, metric, stage, device='cuda'):
        super().__init__(scaler, model, criterion, metric, stage, device, False)

    def on_epoch_start(self):
        self.model.eval()

    @torch.no_grad()
    def batch_update(self, x, y):
        pred = self.model.forward(x)
        if isinstance(self.criterion, nn.Module):
            loss = self.criterion(pred, y)
        else:
            if self.device == 'cuda':
                loss = self.criterion(pred, y, per_image=False)
            else:
                loss = self.criterion(pred, y, per_image=True)
        if isinstance(pred, list):
            pred = pred[0]
        return loss, pred
